import pytorch_mask_rcnn as pmr
from torch.utils import data
import torch


def main(args):
    device = torch.device("cpu")

    # ---------------------- prepare data loader ------------------------------- #

    dataset_train = pmr.datasets(args.dataset, args.data_dir, "train2017", train=True)
    indices = torch.randperm(len(dataset_train)).tolist()
    d_train = torch.utils.data.Subset(dataset_train, indices)

    # ------------------------prepare model-------------------------------------- #

    model = pmr.maskrcnn_resnet50(True, num_classes=2).to(device)
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(
        params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    start_epoch = 0

    # ------------------------------- train ------------------------------------ #

    for epoch in range(start_epoch, args.epochs):
        print("\nepoch: {}".format(epoch + 1))
        pmr.train_one_epoch(model, optimizer, d_train, device, epoch, args)

    # ------------------------------- validation ------------------------------- #

    pmr.generate_results(model, d_train, device)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument("--dataset", default="coco", help="coco or voc")
    parser.add_argument("--data-dir", default="coco2017")
    parser.add_argument("--lr", type=float, default=0.0001)
    parser.add_argument("--momentum", type=float, default=0.9)
    parser.add_argument("--weight_decay", type=float, default=0.0001)
    parser.add_argument("--epochs", type=int, default=500)
    args = parser.parse_args()

    if args.lr is None:
        args.lr = 0.02 * 1 / 16  # lr should be 'batch_size / 16 * 0.02'

    main(args)
